from datetime import datetime
import os
from PyQt5.QtCore import Qt
from PyQt5.QtWidgets import QApplication, QDialog
from dialog import Ui_Dialog
from PyQt5.QtGui import QIcon

from utils.execute_command import execute_in_terminal
from utils.popup_message import popup_info_box, popup_error_box

def translate_to_abs_path(path):
    return os.path.abspath(os.path.expanduser(path))

class MyDialog(Ui_Dialog, QDialog):
    def __init__(self):
        super().__init__()
        self.setupUi(self)
        self.show()

    def setupUi(self, Dialog):
        super().setupUi(Dialog)

        self.setWindowTitle('🦆 Duck Assistant 🦆')
        self.setWindowFlags(Qt.Window)
        self.setWindowIcon(QIcon('res/images/duck_assistant_nobg.png')) # does not work in Ubuntu24, don't know why yet :(

        self.txtPlaygroundDir.setText('~/Projects/Open_Duck_Mini/Open_Duck_Playground')
        self.btnBrowsePlaygroundDir.clicked.connect(self.onBrowsePlaygroundDir)

        self.rdoTrainFromZero.clicked.connect(self.onPolicyFromZeroClicked)
        self.rdoTrainFromCheckpoint.clicked.connect(self.onPolicyFromCheckpointClicked)
        self.rdoTrainFromZero.setChecked(True)
        self.onPolicyFromZeroClicked()
        self.txtTrainFromCheckpointDir.setText('~/Projects/Open_Duck_Mini/Open_Duck_Playground/checkpoints/') # todo, should load from history
        self.btnBrowseCheckpointDir.clicked.connect(self.onBrowseCheckpointDir)
        self.txtTrainSteps.setText('150000000')
        self.btnTrainStart.clicked.connect(self.onStartTrain)
        
        self.txtInferExpDir.setText('~/Projects/Open_Duck_Mini/Open_Duck_Playground/checkpoints/') # todo, should load from history
        self.btnInferBrowseExpDir.clicked.connect(self.onInferBrowseExpDir)
        self.btnInferReloadExpFiles.clicked.connect(self.onInferReloadExpOnnxFiles)
        self.btnInferSaveSummary.clicked.connect(self.onInferSaveSummary)
        self.txtInferExpSummary.textChanged.connect(self.on_infer_exp_summary_changed)
        self.btnInferStartSim.clicked.connect(self.onInferStartSim)
        self.btnInferStartReal.clicked.connect(self.onInferStartReal)

        self.update_onnx_files(pop_error_when_dir_not_exist=False)
        self.load_exp_memo_summary(pop_error_when_dir_not_exist=False)

    def keyPressEvent(self, event):
        if event.key() == Qt.Key_Escape:
            event.ignore()  # 忽略ESC键
        # 捕获 Ctrl+S
        elif event.key() == Qt.Key_S and event.modifiers() == Qt.ControlModifier:
            if self.btnInferSaveSummary.isEnabled():
                self.onInferSaveSummary()
            event.accept()  
        else:
            super().keyPressEvent(event)


    """
    基础设置 ---------------------------------------------------------------------------------------------------------------
    """
    def onPolicyFromZeroClicked(self):
        self.txtTrainFromCheckpointDir.setEnabled(False)
        self.btnBrowseCheckpointDir.setEnabled(False)

    def onPolicyFromCheckpointClicked(self):
        self.txtTrainFromCheckpointDir.setEnabled(True)
        self.btnBrowseCheckpointDir.setEnabled(True)

    def onBrowsePlaygroundDir(self):
        from PyQt5.QtWidgets import QFileDialog
        dir = translate_to_abs_path(self.txtPlaygroundDir.text())
        playground_dir = QFileDialog.getExistingDirectory(self, '选择 Open_Duck_Playground 仓库所在的文件夹...', dir)
        if playground_dir:
            self.txtPlaygroundDir.setText(playground_dir)


    """
    训练 ---------------------------------------------------------------------------------------------------------------
    """
    def onBrowseCheckpointDir(self):
        from PyQt5.QtWidgets import QFileDialog
        dir = translate_to_abs_path(self.txtTrainFromCheckpointDir.text())
        checkpoint_dir = QFileDialog.getExistingDirectory(self, '选择 策略检查点 所在的文件夹...', dir)
        if checkpoint_dir:
            self.txtTrainFromCheckpointDir.setText(checkpoint_dir)

    def onStartTrain(self):
        try:
            main_command = 'uv run playground/open_duck_mini_v2/runner.py'

            # 基座模型
            if self.rdoTrainFromZero.isChecked():
                pass
            elif self.rdoTrainFromCheckpoint.isChecked():
                if not os.path.exists(translate_to_abs_path(self.txtTrainFromCheckpointDir.text())):
                    popup_error_box('错误', '检查点文件夹不存在！')
                    return
                
                main_command += f' --restore_checkpoint_path \'{translate_to_abs_path(self.txtTrainFromCheckpointDir.text())}\''
            else:
                assert False

            # 输出文件夹
            output_dir = os.path.join(translate_to_abs_path(self.txtPlaygroundDir.text()), 'checkpoints')
            main_command += f' --output_dir \'{output_dir}\''

            # 实验名称
            if self.txtTrainExpName.text() != '':
                exp_name = f'{datetime.now().strftime("%m%d_%H%M%S")} - {self.txtTrainExpName.text()}'
            else:
                exp_name = f'{datetime.now().strftime("%m%d_%H%M%S")}'

            main_command += f' --exp_name \'{exp_name}\''

            # 训练步数
            if self.txtTrainSteps.text() != '':
                if not self.txtTrainSteps.text().isdigit():
                    popup_error_box('错误', '训练步数必须为整数！')
                    return
                
                main_command += f' --num_timesteps {self.txtTrainSteps.text()}'

            # 额外参数
            if self.txtTrainExtraParams.text() != '':
                main_command += f' {self.txtTrainExtraParams.text()}'

            # 实验备忘录
            exp_dir = os.path.join(output_dir, exp_name)
            if os.path.exists(exp_dir):
                popup_error_box('错误', '实验目录已存在！')
                return
            
            os.makedirs(exp_dir, exist_ok=True)
            with open(os.path.join(exp_dir, 'exp_memo.txt'), 'w') as f:
                f.write(f'实验名称: {exp_name}\n')

                if self.rdoTrainFromZero.isChecked():
                    f.write('本次训练从零开始训练\n')
                elif self.rdoTrainFromCheckpoint.isChecked():
                    f.write(f'本次训练从 {translate_to_abs_path(self.txtTrainFromCheckpointDir.text())} 中恢复训练\n')

                f.write(f'训练步数: {self.txtTrainSteps.text()}\n')
                f.write(f'额外参数: {self.txtTrainExtraParams.text()}\n')
                f.write('\n')

                f.write('实验备忘信息:\n')
                f.write(self.txtTrainExpMemo.toPlainText())

            # 执行
            execute_in_terminal([
                    f'cd \'{translate_to_abs_path(self.txtPlaygroundDir.text())}\'', 
                    main_command, 
                    ], echo_each_command=True, keep_terminal_open_on_exit=True)
        except Exception as e:
            popup_error_box('错误', str(e))


    """
    推理 ---------------------------------------------------------------------------------------------------------------
    """
    def onInferBrowseExpDir(self):
        from PyQt5.QtWidgets import QFileDialog
        dir = translate_to_abs_path(self.txtInferExpDir.text())
        exp_dir = QFileDialog.getExistingDirectory(self, '选择实验文件夹...', dir)
        if exp_dir:
            self.txtInferExpDir.setText(exp_dir)
            self.update_onnx_files(pop_error_when_dir_not_exist=True)
            self.load_exp_memo_summary(pop_error_when_dir_not_exist=True)

    def onInferReloadExpOnnxFiles(self):
        dir = translate_to_abs_path(self.txtInferExpDir.text())
        if not os.path.exists(dir) or not os.path.isdir(dir):
            popup_error_box('错误', '实验文件夹不存在！')
            return

        self.update_onnx_files(pop_error_when_dir_not_exist=True)
        self.load_exp_memo_summary(pop_error_when_dir_not_exist=True)

    def update_onnx_files(self, pop_error_when_dir_not_exist):
        self.lstInferOnnxFiles.clear()

        dir = translate_to_abs_path(self.txtInferExpDir.text())
        if not os.path.exists(dir) or not os.path.isdir(dir): 
            if pop_error_when_dir_not_exist:
                popup_error_box('错误', '实验文件夹不存在！')

            return

        onnx_files = [f for f in os.listdir(dir) if f.endswith('.onnx')]
        onnx_files.sort()
        self.lstInferOnnxFiles.addItems(onnx_files)

    def load_exp_memo_summary(self, pop_error_when_dir_not_exist):
        dir = translate_to_abs_path(self.txtInferExpDir.text())
        if not os.path.exists(dir) or not os.path.isdir(dir): 
            if pop_error_when_dir_not_exist:
                popup_error_box('错误', '实验文件夹不存在！')

            return

        # exp_memo
        exp_memo_path = os.path.join(dir, 'exp_memo.txt')
        if os.path.exists(exp_memo_path):
            with open(exp_memo_path, 'r') as f:
                self.txtInferExpMemo.setPlainText(f.read())
        else:
            self.txtInferExpMemo.setPlainText('')

        # exp_summary
        exp_summary_path = os.path.join(dir, 'exp_summary.txt')
        if os.path.exists(exp_summary_path):
            with open(exp_summary_path, 'r') as f:
                self.txtInferExpSummary.setPlainText(f.read())
        else:
            self.txtInferExpSummary.setPlainText('')

        self.btnInferSaveSummary.setEnabled(False)

    def onInferSaveSummary(self):
        dir = translate_to_abs_path(self.txtInferExpDir.text())
        if not os.path.exists(dir) or not os.path.isdir(dir): 
            popup_error_box('错误', '实验文件夹不存在！')
            return

        exp_summary_path = os.path.join(dir, 'exp_summary.txt')
        with open(exp_summary_path, 'w') as f:
            f.write(self.txtInferExpSummary.toPlainText())

        self.btnInferSaveSummary.setEnabled(False)

    def on_infer_exp_summary_changed(self):
        self.btnInferSaveSummary.setEnabled(True)

    def onInferStartSim(self):
        if self.lstInferOnnxFiles.currentItem() is None:
            popup_error_box('错误', '请选择一个 ONNX 文件来推理！')
            return
        
        try:
            playground_dir = translate_to_abs_path(self.txtPlaygroundDir.text())
            onnx_path = os.path.join(translate_to_abs_path(self.txtInferExpDir.text()), self.lstInferOnnxFiles.currentItem().text())

            execute_in_terminal([
                f'cd \'{playground_dir}\'', 
                f'(uv run playground/open_duck_mini_v2/mujoco_infer.py --onnx_model_path \'{onnx_path}\' || true)'
                ], echo_each_command=True, keep_terminal_open_on_exit=False)
        except Exception as e:
            popup_error_box('错误', str(e))

    def onInferStartReal(self):
        if self.lstInferOnnxFiles.currentItem() is None:
            popup_error_box('错误', '请选择一个 ONNX 文件来推理！')
            return
        
        try:
            onnx_path = os.path.join(translate_to_abs_path(self.txtInferExpDir.text()), self.lstInferOnnxFiles.currentItem().text())

            execute_in_terminal([
                f'echo -e \'\\e[1;32m传输 {onnx_path} 到真机...\\e[0m\'', 
                f'sshpass -p 1936 scp -p \'{onnx_path}\' richard@duck.local:~/Open_Duck_Mini_Runtime/remote_walk/temp.onnx', 
                f'echo -e \'\\e[1;32m执行行走脚本...\\e[0m\'',
                f'sshpass -p 1936 ssh -t richard@duck.local ~/Open_Duck_Mini_Runtime/remote_walk/remote_walk.sh'
                ], echo_each_command=False, keep_terminal_open_on_exit=False)
        except Exception as e:
            popup_error_box('错误', str(e))



if __name__ == '__main__':
    app = QApplication([])
    dialog = MyDialog()
    app.exec_()
